import pandas as p
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import sys

#Plot ticks in middle of plot
def plotTick(ax, xLoc, shift=False):
    height = .05
    zorder = 1
    ax.plot([xLoc, xLoc], [height, -height], color='xkcd:black', zorder=zorder)
    if(shift and xLoc>1):
        ax.text(xLoc, -height*1.5, str(int(xLoc)+8), ha='center', va='top', zorder=zorder)
    else:
        ax.text(xLoc, -height * 1.5, str(int(xLoc)), ha='center', va='top', zorder=zorder)
#Plot multiple ticks
def plotTicks(ax, xLocs, shift=False):
    xlim = ax.get_xlim()
    minXLim = min(xlim)
    maxXlim = max(xlim)
    for xLoc in xLocs:
        if(xLoc>minXLim and xLoc<maxXlim):
            plotTick(ax, xLoc, shift)

if(sys.platform=='darwin'):
    file = 'Tables/iv.xlsx'
else:
    file = r'Tables\iv.xlsx'

dataDict = p.read_excel(file, None, None)

vmtData = dataDict['Spec2']
gasUseData = dataDict['Spec2GasUse']

confSetRowName = '95% Confidence Set'
confSetIndVMT = vmtData.iloc[:, 0] == confSetRowName
confSetIndGasUse = gasUseData.iloc[:, 0] == confSetRowName

betaVMTCS = float(vmtData.loc[confSetIndVMT, 1]), float(vmtData.loc[confSetIndVMT, 2])
betaGasUseCS = float(gasUseData.loc[confSetIndGasUse, 1]), float(gasUseData.loc[confSetIndGasUse, 2])

betaRowName = 'beta'
betaIndVMT = vmtData.iloc[:, 0] == betaRowName
betaIndGasUse = gasUseData.iloc[:, 0] == betaRowName

betaVMT = float(vmtData.loc[betaIndVMT, 1])
betaGasUse = float(gasUseData.loc[betaIndGasUse, 1])

jacobsenVMT = 1
jacobsenGasUse = 1

priaVMT = 1.6
priaGasUse = -4.8

priaStarVMT = 0.3
priaStarGasUse = -6.1

friaVMT = -2.5
friaGasUse = 9.7

friaStarVMT = 0
friaStarGasUse = 11

offset = 3
nullColors = sns.color_palette('bright', n_colors=offset+9)[offset:]
jacobsenColor = nullColors[0]
priaColor = nullColors[1]
friaColor = nullColors[2]
priaStarColor = nullColors[3]
friaStarColor = nullColors[7]

supTitleHeight = 1
dpi = 300


#VMT Plot
x = 6
gr = (1+np.sqrt(5))/2
y = x/(gr)
dotSize = 75
lineWidth = 3

elasticitiesDict = {'Elast. Lower Bound': min(betaVMTCS), 'Elast. Upper Bound': max(betaVMTCS), 'Est. Elast.':betaVMT, 'Jacobsen':jacobsenVMT, 'PRIA':priaVMT, 'FRIA':friaVMT, 'PRIA-Star':priaStarVMT, 'FRIA-Star':friaStarVMT}

fig, lax = plt.subplots(1,1,figsize=(x,y))

lax.text(elasticitiesDict['Elast. Lower Bound'], 1 / 2+.1, '95% Confidence Set', ha='center', va='bottom', color='xkcd:green')
# lax.text(elasticitiesDict['Elast. Upper Bound'], 1 / 2+.1, '95% Confidence\nUpper Bound', ha='center', va='bottom', color='xkcd:green')

lax.text(elasticitiesDict['Est. Elast.']-.5, .5 / 2, 'Estimated\nElasticity', ha='center', va='center', color='xkcd:dark green')

lax.text(elasticitiesDict['Jacobsen'], -.6 / 2, 'Unit\nElasticity', ha='left', va='center', color=jacobsenColor)

lax.text(elasticitiesDict['PRIA'], .4 / 2, 'SAFE\nPRIA', ha='center', va='center', color=priaColor)
lax.text(elasticitiesDict['PRIA-Star']+.25, +.5 / 2, 'SAFE\nPRIA*', ha='center', va='center', color=priaStarColor)

lax.text(elasticitiesDict['FRIA'], .4 / 2, 'SAFE\nFRIA', ha='center', va='center', color=friaColor)
lax.text(elasticitiesDict['FRIA-Star'], -.6 / 2, 'SAFE\nFRIA*', ha='center', va='center', color=friaStarColor)

xlim = plt.xlim()


lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound'] - .25], [-1 / 2, -1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound'] - .25], [1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)

lax.scatter([elasticitiesDict['Est. Elast.'], elasticitiesDict['Est. Elast.']], [0, 0], color='xkcd:dark green', s=dotSize + 20, zorder=3, marker='d')

lax.scatter([elasticitiesDict['Jacobsen'], elasticitiesDict['Jacobsen']], [0, 0], color=jacobsenColor, s=dotSize, zorder=2)

lax.scatter([elasticitiesDict['PRIA'], elasticitiesDict['PRIA']], [0, 0], color=priaColor, marker='o', s=dotSize, zorder=2)

lax.scatter([elasticitiesDict['FRIA'], elasticitiesDict['FRIA']], [0, 0], color=friaColor, marker='o', s=dotSize, zorder=3)

lax.scatter([elasticitiesDict['PRIA-Star'], elasticitiesDict['PRIA-Star']], [0, 0], color=priaStarColor, marker='o', s=dotSize, zorder=5)

lax.scatter([elasticitiesDict['FRIA-Star'], elasticitiesDict['FRIA-Star']], [0, 0], color=friaStarColor, marker='o', s=dotSize, zorder=5)



lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Upper Bound']], [0, 0], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Upper Bound']], [0, 0], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound'] + .25], [-1 / 2, -1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound'] + .25], [1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)

lax.plot([-800, 350], [0,0], color='xkcd:black', zorder=1)


rightLims = [8, 9.5]
leftLims = [-6.5, 3.81]

lax.set_xlim(leftLims)

leftXTicks = lax.get_xticks()
newLeftXTicks = leftXTicks[:len(leftXTicks)-1]
lax.set_xticks(newLeftXTicks)

plt.ylim([-1,1])

plt.yticks([])

plotTicks(lax, range(-6, 4))


lax.set_xticks([])
plt.suptitle('Fleet Distance Traveled Elasticity\nEstimate and Null Hypotheses',y=supTitleHeight)

xlim = lax.get_xlim()
ylim = lax.get_ylim()

lax.fill_between([min(xlim), 0], max(ylim), min(ylim), color='xkcd:gray', alpha=.25)
lax.set_ylim(ylim)


plt.xlim((-6.5, 3.81))


if(sys.platform=='darwin'):
    saveFile = 'PythonScripts/Plots/Elasticities/VMTElasticities.png'
else:
    saveFile = r'PythonScripts\Plots\Elasticities\VMTElasticities.png'

plt.savefig(saveFile, bbox_inches='tight', dpi=dpi)

#Gas Use Plot

elasticitiesDict = {'Elast. Lower Bound': min(betaGasUseCS), 'Elast. Upper Bound': max(betaGasUseCS), 'Est. Elast.':betaGasUse, 'Jacobsen':jacobsenGasUse, 'PRIA':priaGasUse, 'FRIA':friaGasUse, 'PRIA-Star':priaStarGasUse, 'FRIA-Star':friaStarGasUse}


fig, lax = plt.subplots(1,1,figsize=(x,y))


lax.text(elasticitiesDict['Elast. Lower Bound'], 1 / 2+.1, '95% Confidence\nConfidence Set', ha='left', va='bottom', color='xkcd:green')
lax.text(elasticitiesDict['Est. Elast.']-.5, .5 / 2, 'Estimated\nElasticity', ha='center', va='center', color='xkcd:dark green')

lax.text(elasticitiesDict['Jacobsen'], -.6 / 2, 'Unit\nElasticity', ha='left', va='center', color=jacobsenColor)

lax.text(elasticitiesDict['PRIA'], .4 / 2, 'SAFE\nPRIA', ha='center', va='center', color=priaColor)

lax.text(elasticitiesDict['FRIA']-8, .4 / 2, 'SAFE\nFRIA', ha='center', va='center', color=friaColor)


lax.text(elasticitiesDict['PRIA-Star']+.5, -.6 / 2, 'SAFE\nPRIA*', ha='center', va='center', color=priaStarColor)

lax.text(elasticitiesDict['FRIA-Star']-8, -.6 / 2, 'SAFE\nFRIA*', ha='center', va='center', color=friaStarColor)

xlim = plt.xlim()


lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound'] - .25], [-1 / 2, -1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Upper Bound'], elasticitiesDict['Elast. Upper Bound'] - .25], [1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)

lax.scatter([elasticitiesDict['Est. Elast.'], elasticitiesDict['Est. Elast.']], [0, 0], color='xkcd:dark green', s=dotSize + 20, zorder=3, marker='d')

lax.scatter([elasticitiesDict['Jacobsen'], elasticitiesDict['Jacobsen']], [0, 0], color=jacobsenColor, s=dotSize, zorder=2)

lax.scatter([elasticitiesDict['PRIA'], elasticitiesDict['PRIA']], [0, 0], color=priaColor, marker='o', s=dotSize, zorder=2)

lax.scatter([elasticitiesDict['FRIA']-8, elasticitiesDict['FRIA']-8], [0, 0], color=friaColor, marker='o', s=dotSize, zorder=3)

lax.scatter([elasticitiesDict['PRIA-Star'], elasticitiesDict['PRIA-Star']], [0, 0], color=priaStarColor, marker='o', s=dotSize, zorder=5)

lax.scatter([elasticitiesDict['FRIA-Star']-8, elasticitiesDict['FRIA-Star']-8], [0, 0], color=friaStarColor, marker='o', s=dotSize, zorder=5)



lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Upper Bound']], [0, 0], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Upper Bound']], [0, 0], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound']], [-1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound'] + .25], [-1 / 2, -1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)
lax.plot([elasticitiesDict['Elast. Lower Bound'], elasticitiesDict['Elast. Lower Bound'] + .25], [1 / 2, 1 / 2], color='xkcd:green', linewidth=lineWidth, zorder=2)

lax.plot([-800, 350], [0,0], color='xkcd:black', zorder=1)


rightLims = [8, 13]
leftLims = [-6.5, 2.5]
leftLims = [-6.5, 3.81]

lax.scatter(1.36, 0, marker='s', color='xkcd:white', zorder=5, s=50, edgecolor='xkcd:white')
lax.plot([1.45, 1.55], [-.05, .05], color='xkcd:black', zorder=6)
lax.plot([1.18, 1.28], [-.05, .05], color='xkcd:black', zorder=6)





lax.set_xlim(leftLims)

leftXTicks = lax.get_xticks()
newLeftXTicks = leftXTicks[:len(leftXTicks)-1]
lax.set_xticks(newLeftXTicks)

plt.ylim([-1,1])


plotTicks(lax, range(-6, 4), shift=True)


lax.set_xticks([])
lax.set_yticks([])
plt.suptitle('Highway Gas Use Elasticity\nEstimate and Null Hypotheses',y=supTitleHeight)

plt.xlim((-6.5, 3.81))
xlim = lax.get_xlim()
ylim = lax.get_ylim()

lax.fill_between([min(xlim), 0], max(ylim), min(ylim), color='xkcd:gray', alpha=.25)
lax.set_ylim(ylim)


if(sys.platform=='darwin'):
    saveFile = 'PythonScripts/Plots/Elasticities/GasUseElasticities.png'
else:
    saveFile = r'PythonScripts\Plots\Elasticities\GasUseElasticities.png'


plt.savefig(saveFile, bbox_inches='tight', dpi=dpi)